import os
import torch
import SimpleITK as sitk
import matplotlib.pyplot as plt

"""
    通过检测像素值是否存在来得到海马体的大致区间(横断面)
"""
# 标签文件夹地址
mask_path = 'C:/SYBdataset/dataset/image_registration/AffinedManualSegImageNIfTI'
# mask_path = 'C:/SYBdataset/dataset/AffinedManualSegImageNIfTI'
# 获得标签地址列表
mask_path_files = [os.path.join(mask_path, f) for f in os.listdir(mask_path) if f.endswith('.nii.gz')]
mask_max = 0  # 横断面区间最大值
mask_min = 256  # 横断面区间最小值
flag = 0  # 标志位,用于判断是否记录过最小值
for i in range(len(mask_path_files)):  # 扫描图像横截面
    # 读取图片
    img = sitk.ReadImage(mask_path_files[i])
    print(i)
    img_np = sitk.GetArrayFromImage(img)  # (176, 256, 224)
    mask = torch.from_numpy(img_np).long()
    mask = mask.permute(1, 2, 0)  # 换轴(176, 256, 224)->(256, 224, 176)
    for j in range(mask.shape[0]):
        if mask[j].sum() >= 5 and flag == 0 and j < mask_min:
            mask_min = j  # 记录最小值
            print('mask_min', mask_min)
            print('mask[j].sum()_min', mask[j].sum())
            flag = 1  # 已经记录过最小值
            # 显示此时最小值的横截面图像
            plt.subplots_adjust(top=1)
            slice_1 = mask[j]
            im1 = plt.imshow(slice_1, cmap='gray')  # 使用灰度映射显示图片
            cbar1 = plt.colorbar(im1)  # 加个灰度的颜色条
            plt.show()

        if mask[j].sum() <= 5 and flag == 1 and j > mask_max:
            mask_max = j
            print('mask_max', mask_max)
            print('mask[j].sum()_max', mask[j].sum())
            flag = 0  # 已经记录过最大值,清零
            break
print('mask_max', mask_max)  # 输出最大值
print('mask_min', mask_min)  # 输出最小值
